import torch
import flair
import time
from flair.data import Sentence
from flair.embeddings import WordEmbeddings, XLNetEmbeddings, BertEmbeddings, DocumentPoolEmbeddings, \
    DocumentRNNEmbeddings, TransformerWordEmbeddings, TransformerDocumentEmbeddings
import eco_prj.representation.rere_config as cnf


class FlairEmbedder():

    def __init__(self, method='bert'):
        super(FlairEmbedder, self).__init__()
        flair.device = cnf.device
        if method == 'bert':
            self.embedder = TransformerWordEmbeddings('roberta-base')
        else:
            self.embedder = WordEmbeddings('glove')
        self.document_embedder = DocumentPoolEmbeddings([self.embedder])
        self.embedding_dim = self.embedder.embedding_length

    def embed(self, sentence: str):
        # print(sentence)
        vecs = []
        sen = Sentence(sentence)
        self.embedder.embed(sen)
        for token in sen:
            if token.text != '.':
                vecs.append(token.embedding)
        vecs = torch.stack(vecs)
        return vecs

    def embed_chunks(self, chunks_sentence: str, separator=':::'):
        """

        :param chunk_sentence:
        :param separator:
        :return:
        """
        chunks = chunks_sentence.split(separator)
        vecs = []
        a = None
        try:
            for sentence in chunks:
                sen = Sentence(sentence)
                self.document_embedder.embed(sen)
                vec = sen.get_embedding()
                vecs.append(vec)
            # out_matrix = torch.Tensor(len(chunks), self.embedding_dim)
            a = torch.stack(vecs)
        except Exception as e:
            print(e)
            print(chunks_sentence)
        return a

    def embed_a_chunk(self, chunk: str):
        """

        :param chunk_sentence:
        :param separator:
        :return:
        """
        vec = None
        try:
            sen = Sentence(chunk)
            self.document_embedder.embed(sen)
            vec = sen.get_embedding()
        except Exception as e:
            print(e)
            print(chunk)
        return vec


if __name__ == '__main__':
    model = FlairEmbedder('glove')
    s = 'Red, white and bleu salad:::super yum:::a great addition:::the menu:::This location:::great service:::food:::just the right temps:::Kids pizza:::a hit:::lots:::great side dish options:::the kiddos:::town:::a spot'
    s = '环境监测数据认可的申请'
    start = time.time()
    for i in range(1):
        # vs = model.embed_chunk(s)
        vs = model.embed(s)
        print(vs)
    end = time.time()
    print(end - start)
